token_data = {
    "access_token": credentials.token,
    "token_type": "Bearer",
    "expires_at": credentials.expiry.timestamp(),  # Unix timestamp
    "scope": "https://www.googleapis.com/auth/cloud-platform",
    "cached_at": time.time(),  # For debugging
    "provider": "vertex"  # Helpful for multi-provider systems
}

async_oauth_token_cache: "AsyncCache" = _AsyncBackend(ttl_seconds=None)

async def get_cached_oauth_token_async(api_key: str) -> dict[str, Any] | None:
    if not api_key:
        return None
    import hashlib, time
    cache_key = f"token:{hashlib.sha256(api_key.encode()).hexdigest()}"
    cached_data = await async_oauth_token_cache.get(cache_key)
    if not cached_data:
        return None
    expires_at = cached_data.get("expires_at")
    if not expires_at:
        await async_oauth_token_cache.delete(cache_key)
        return None
    current_time = time.time()
    if expires_at <= current_time:
        await async_oauth_token_cache.delete(cache_key)
        await _opportunistic_cleanup(current_time, max_items=2)
        return None
    return cached_data

async def _opportunistic_cleanup(current_time: float, max_items: int = 2):
    cleaned = 0
    if hasattr(async_oauth_token_cache, "cache"):
        for key, value in list(async_oauth_token_cache.cache.items()):
            if cleaned >= max_items:
                break
            if key.startswith("token:"):
                expires_at = value.get("expires_at")
                if expires_at and expires_at <= current_time:
                    await async_oauth_token_cache.delete(key)
                    cleaned += 1
    elif hasattr(async_oauth_token_cache, "client"):
        try:
            import os
            pattern = f"{os.getenv('REDIS_PREFIX', 'forge')}:token:*"
            async for redis_key in async_oauth_token_cache.client.scan_iter(match=pattern, count=10):
                if cleaned >= max_items:
                    break
                key_str = redis_key.decode() if isinstance(redis_key, bytes) else redis_key
                internal_key = key_str.split(":", 1)[-1]
                cached_data = await async_oauth_token_cache.get(internal_key)
                if cached_data:
                    expires_at = cached_data.get("expires_at")
                    if expires_at and expires_at <= current_time:
                        await async_oauth_token_cache.delete(internal_key)
                        cleaned += 1
        except Exception:
            pass

async def cache_oauth_token_async(api_key: str, token_data: dict[str, Any]) -> None:
    if not api_key or not token_data:
        return
    import hashlib
    cache_key = f"token:{hashlib.sha256(api_key.encode()).hexdigest()}"
    if "expires_at" not in token_data:
        logger.warning("OAuth token cached without expires_at - skipping")
        return
    await async_oauth_token_cache.set(cache_key, token_data)

def get_vertex_token(api_key):
    token_data = redis.get(f"vertex_token:{api_key}")
    if not token_data:
        raise Unauthorized("Vertex not bound")

    # if the access token is not expired
    if token_data["expires_at"] > current_timestamp():
        return token_data["access_token"]

    # if there has refresh_token, then refresh
    if "refresh_token" in token_data:
        new_token = refresh_vertex_token(token_data["refresh_token"])
        redis.set(f"vertex_token:{api_key}", new_token, ttl=3600)
        return new_token["access_token"]

    # can not refresh, not to re-login
    raise TokenExpired("Please reconnect your Vertex account")